import asyncio
from enum import Enum
from typing import Dict, Any

from autogen_core import CancellationToken
from autogen_ext.tools.mcp import StdioServerParams
from pydantic import BaseModel

from mcp_tools.BaseMcpClient import BaseMCPClient


class DatabaseType(Enum):
    """数据库类型枚举"""
    MYSQL = "mysql"
    POSTGRESQL = "postgresql"
    SQLITE = "sqlite"
    ORACLE = "oracle"
    SQLSERVER = "sqlserver"


class DummyParams(BaseModel):
    sql: str


class DataBaseMcp(BaseMCPClient):
    """
    数据库MCP
    """

    def __init__(self, db_type: DatabaseType = DatabaseType.MYSQL):
        super().__init__()
        self.db_type = db_type

    def get_params(self) -> StdioServerParams:
        return StdioServerParams(
            command="npx",
            args=["-y", "@bytebase/dbhub", "--transport", "stdio", "--readonly", "--dsn",
                  "mysql://root:123456@localhost:3306/continew_admin?sslmode=disable"],
            read_timeout_seconds=60
        )

    def gettype(self) -> str:
        """获取数据库类型"""
        return self.db_type.value

    def get_sql_syntax_info(self) -> Dict[str, Any]:
        """获取SQL语法信息"""
        if self.db_type == DatabaseType.MYSQL:
            return {
                "type": "mysql",
                "date_format": "YYYY-MM-DD",
                "datetime_format": "YYYY-MM-DD HH:MM:SS",
                "string_quote": "'",
                "identifier_quote": "`",
                "limit_syntax": "LIMIT",
                "offset_syntax": "OFFSET"
            }
        elif self.db_type == DatabaseType.POSTGRESQL:
            return {
                "type": "postgresql",
                "date_format": "YYYY-MM-DD",
                "datetime_format": "YYYY-MM-DD HH:MM:SS",
                "string_quote": "'",
                "identifier_quote": "\"",
                "limit_syntax": "LIMIT",
                "offset_syntax": "OFFSET"
            }
        elif self.db_type == DatabaseType.SQLITE:
            return {
                "type": "sqlite",
                "date_format": "YYYY-MM-DD",
                "datetime_format": "YYYY-MM-DD HH:MM:SS",
                "string_quote": "'",
                "identifier_quote": "\"",
                "limit_syntax": "LIMIT",
                "offset_syntax": "OFFSET"
            }
        else:
            return {
                "type": "unknown",
                "date_format": "YYYY-MM-DD",
                "datetime_format": "YYYY-MM-DD HH:MM:SS",
                "string_quote": "'",
                "identifier_quote": "\"",
                "limit_syntax": "LIMIT",
                "offset_syntax": "OFFSET"
            }


async def test_mcp():
    db_mcp = DataBaseMcp()
    tools = await db_mcp.init_tools()  # 初始化工具，启动 MCP 服务

    # tools 是一个列表，里面是 MCP 封装的工具函数
    print(f"已加载 MCP 工具数量: {len(tools)}")
    for t in tools:
        print(f"工具名: {t.name}, 描述: {t.description}")

    # 测试调用第一个 MCP 工具，传入示例参数
    tool = tools[0]
    # 这里用 asyncio.run 调用异步的 MCP 工具方法
    # 你可以根据工具定义调整参数内容
    cancellation_token = CancellationToken()
    params = DummyParams(sql="select * from sys_user;")

    try:
        # 用字典传入参数，注意参数名根据具体MCP接口调整
        result = await tool.run(params, cancellation_token)
        print(dir(result))
        print("调用结果：", result[0].text)
    except Exception as e:
        import traceback
        print("调用异常：")
        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(test_mcp())
